/* merge.c
 * This file is part of Landmerge (http://code.google.com/p/landmerge/)
 * Copyright (c) Matus Tejiscak 2008
 * Covered by GNU GPL v3 (http://www.gnu.org/licenses/gpl.txt).
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <math.h>
#include <string.h>

#include "tiff.h"

// Symbolic indexes to various structures
#define R 0
#define G 1
#define B 2
#define I 3
#define P 4

/* ### TYPE DEFINITIONS ### */

/// Set of input TIFFs.
typedef struct tiff_in_t
{
	tiff_t *chan[5];
}
tiff_in_t;

/// Image statistics. (Applies on grayscale images only).
typedef struct stats_t
{
	/// \sum x_i
	double Sx;
	/// avg(x_i)
	double avg;
	/// \sum 1
	unsigned long pixel_count;

	/// \sum (x_i - avg_x)^2
	double Sdd;
	/// standard deviation
	double sd;
}
stats_t;

/// Stats used for linear regression.
typedef struct regr_stats_t
{
	double Sxx, Sxy, Sxz, Syz, Syy, Szz;
	double Sx, Sy, Sz, Sb;
	double Sbx, Sby, Sbz;
}
regr_stats_t;

/// Actions available in the main loop.
typedef enum action_t
{
	AC_REGRESSION,
	AC_STATS_1,
	AC_STATS_2,
	AC_MERGE
}
action_t;

const char * const str_action[] = {
	"calculating linear regression coefficients...",
	"gathering statistics (stage1)...",
	"gathering statistics (stage2)...",
	"merging..."
};

/* ### CONSTANTS ### */

// Constants		  R     G     B     IR    Pan
/// Postprocess gamma correction.
//double post_gamma[] =	{ 1.2,  1.0,  1.5 };
double post_gamma[] =	{ 1.15/2.2, 1.10/2.2, 1.33/2.2 };
/// <0..1> pixel range will be mapped to <avg - dyn_range*SD, avg + dyn_range*SD>
/// where SD is standard deviation of the samples.
double dyn_range[] =	{ 5.0,  5.0,  5.0 };
double dyn_range_Y = 3.0;
double brightness = 0.75;
int do_postprocess = 1;

/// Channel subsampling. Constant, according to the Landsat docs.
int subsamp[] =		{ 2,    2,    2,    2,    1 };

int ofsx[] =		{ 0,    0,    0,    0,    0 };
int ofsy[] =            { 0,    0,    0,    0,    0 };

/* ### MUTABLE GLOBAL STATE ### */

// Crop bounds
int crop_x = 2500;
int crop_y = 6200;
int crop_w = 1000;
int crop_h = 1000;

double Sy = 0;
double Sdd = 0;
double avg_Y = 0;
double sd_Y = 0;
unsigned long pixel_count = 0;

/// Statistics
stats_t stats[3];
const stats_t zero_stats = {
	.Sx = 0, .avg = 0, .pixel_count = 0,
	.Sdd = 0, .sd = 0
};

const regr_stats_t zero_regr_stats = {
	.Sxx = 0, .Syy = 0, .Szz = 0, .Sxy = 0, .Syz = 0, .Sxz = 0,
	.Sbx = 0, .Sby = 0, .Sbz = 0,
	.Sx = 0, .Sy = 0, .Sz = 0, .Sb = 0
};
regr_stats_t regr_stats;

double alpha = 0;
double beta = 0;
double gamma = 0;
double phi = 0;
double eta = 0;
double psi = 0;
double tlx, tly, trx, try, blx, bly, brx, bry;

/* ### FUNCTIONS & LOGIC ### */

/// Print an error message and quit.
void die(const char * const reason, ...)
{
	va_list ap;
	va_start(ap, reason);
	vfprintf(stderr, reason, ap);
	va_end(ap);
	exit(1);
}

/// An utility function to print progress.
void progress(const char * const fmt, ...)
{
	va_list ap;
	va_start(ap, fmt);
	vprintf(fmt, ap);
	va_end(ap);
	fflush(stdout);
}

__attribute__ ((gnu_inline)) inline double sqr(double x)
{
	return x * x;
}

__attribute__ ((gnu_inline)) inline double apply_gamma(double x, double gamma)
{
	return  (x < 1.0e-9) ? 0 : exp(gamma * log(x));
}

void add_regr_stats(regr_stats_t *target, const regr_stats_t *src)
{
	target->Sxx += src->Sxx; target->Syy += src->Syy; target->Szz += src->Szz;
	target->Sxy += src->Sxy; target->Syz += src->Syz; target->Sxz += src->Sxz;
	target->Sbx += src->Sbx; target->Sby += src->Sby; target->Sbz += src->Sbz;
	target->Sx += src->Sx; target->Sy += src->Sy; target->Sz += src->Sz; target->Sb += src->Sb;
}

// Utility macros
#define MIN(a, b) ((a < b) ? (a) : (b))
#define MAX(a, b) ((a > b) ? (a) : (b))
#define BOUND(x) ((x < 0) ? 0 : ((x > 1) ? 1 : x))
#define QUANT(x) lround(BOUND(x) * 255)

/// Calculate one pixel's RGB color
void calculate_pixel(double *out, unsigned char *val, stats_t *stats, int apply_postprocess)
{
	double x[5];
	for (int i = 0; i < 5; ++i)
		x[i] = val[i] / 255.0;

	x[P] -= phi / 255.0;

	double mu = alpha*alpha + beta*beta;
	double sigma = 1 - eta;

	out[R] = ((alpha*alpha + beta*beta*sigma)*x[R] - alpha*beta*eta*x[G] - gamma*beta*eta*x[I] + beta*eta*x[P]) / mu;
	out[G] = (-alpha*beta*eta*x[R] + (beta*beta + alpha*alpha*sigma)*x[G] - gamma*alpha*eta*x[I] + alpha*eta*x[P]) / mu;
	out[B] = x[B];

	// False plasticity injection
	double Y  = + 0.299    * out[R] + 0.587    * out[G] + 0.114    * out[B];
	double Pb = - 0.168736 * out[R] - 0.331264 * out[G] + 0.5      * out[B];
	double Pr = + 0.5      * out[R] - 0.418688 * out[G] - 0.081312 * out[B];

	Y = Y*(1-psi) + psi*x[P];

	out[R] = Y + 1.402 * Pr;
	out[G] = Y - 0.344 * Pb - 0.714 * Pr;
	out[B] = Y + 1.772 * Pb;

	// Apply gamma
	for (int i = 0; i < 3; ++i)
	{
		out[i] = apply_gamma(out[i], post_gamma[i]);
	}

	if (apply_postprocess)
	{
		for (int i = 0; i < 3; ++i)
		{
			out[i] = brightness * (1 + (out[i] - avg_Y) / (dyn_range_Y * sd_Y)) / 2.0;
		}
	}
}

void process_images(const tiff_in_t *tiffs, action_t action, stats_t *stats, FILE *out_rgb)
{
	// Print some progress
	printf("%s\n", str_action[(int) action]);

	// The output size is the panchromatic image's size
	unsigned width = tiffs->chan[P]->row_size;
	unsigned height = tiffs->chan[P]->row_count;

	unsigned outline_size = width * 3;
	unsigned char *outline = (unsigned char *) malloc(outline_size);

	// Initialize scanlines
	unsigned char *next[5], *current[5];
	for (int col = 0; col < 5; ++col)
	{
		unsigned size = tiffs->chan[col]->row_size;

		tiff_seek(tiffs->chan[col], (crop_y + ofsy[col] - 1) / subsamp[col]);

		next[col] = (unsigned char *) tiffs->chan[col]->scanline;
		current[col] = (unsigned char *) malloc(tiffs->chan[col]->row_size);
		memcpy(current[col], next[col], tiffs->chan[col]->row_size);

		tiff_next_scanline(tiffs->chan[col]);
	}

	// Loop over the image
	for (int y = crop_y; y < crop_y + crop_h && y < height; ++y)
	{
		// Determine scene borders
		double bound_left, bound_right;

		// Left border
		if (y > bly)
		{
			int Dx = brx - blx;
			int dy = y - bly;
			double Dy = bry - bly;
			bound_left = blx + dy*Dx/Dy;
		}
		else
		{
			int Dx = blx - tlx;
			int dy = y - tly;
			double Dy = bly - tly;
			bound_left = tlx + dy*Dx/Dy;
		}

		// Right border
		if (y > try)
		{
			int Dx = brx - trx;
			int Dy = bry - try;
			double dy = y - try;
			bound_right = trx + dy*Dx/Dy;
		}
		else
		{
			int Dx = trx - tlx;
			int Dy = try - tly;
			double dy = y - tly;
			bound_right = tlx + dy*Dx/Dy;
		}
		
		// Read next scanline(s) if needed
		for (int col = 0; col < 5; ++col)
		{
			if ((y+ofsy[col]) % subsamp[col] == 0)
			{
				tiffs->chan[col]->scanline = current[col];
				current[col] = next[col];
				next[col] = tiffs->chan[col]->scanline;
				tiff_next_scanline(tiffs->chan[col]);
			}
		}

		// Where the pixel will be written
		unsigned char *pixel = outline + 3*crop_x;
		unsigned char val[5];
		regr_stats_t lrs = zero_regr_stats;
		stats_t ls[3];
		double line_S = 0;
		for (int i = 0; i < 3; ++i) ls[i] = zero_stats;


		// Traverse all pixels
		for (int x = crop_x; x < crop_x + crop_w && x < width; ++x)
		{
			// Skip invalid area
			if (x < bound_left || x > bound_right)
			{
				pixel[0] = pixel[1] = pixel[2] = 0;
				pixel += 3;
				continue;
			}

			int x_sampled[5];
			double p[3];
			for (int col = 0; col < 5; ++col)
			{
				// Linear interpolation: calculate the four pixels
				int sx = x / subsamp[col];

				unsigned char top_left = current[col][sx+ofsx[col]];
				unsigned char top_right = (x < width-1) ? current[col][sx+1+ofsx[col]] : top_left;
				unsigned char bottom_left = next[col][sx+ofsx[col]];
				unsigned char bottom_right = (x < width-1) ? next[col][sx+1+ofsx[col]] : bottom_left;

				int xmod = x % subsamp[col];
				int ymod = y % subsamp[col];

				// Interpolate horizontally
				unsigned char upper = (xmod*top_right + (subsamp[col]-xmod)*top_left) / subsamp[col];
				unsigned char lower = (xmod*bottom_right + (subsamp[col]-xmod)*bottom_left) / subsamp[col];

				// Interpolate vertically into the final value
				val[col] = (ymod*lower + (subsamp[col]-ymod)*upper) / subsamp[col];
			}

			switch (action)
			{
				case AC_MERGE:
					// Calculate & write pixel value
					if (val[0] || val[1] || val[2])
					{
						calculate_pixel(p, val, stats, do_postprocess);
						for (int i = 0; i < 3; ++i) pixel[i] = QUANT(p[i]);
					}
					else
					{
						pixel[0] = pixel[1] = pixel[2] = 0;
					}
				break;

				case AC_REGRESSION:
					if (val[0] || val[1] || val[2])
					{
						lrs.Sxx += val[G]*val[G];
						lrs.Syy += val[R]*val[R];
						lrs.Szz += val[I]*val[I];
						lrs.Sxy += val[G]*val[R];
						lrs.Syz += val[R]*val[I];
						lrs.Sxz += val[G]*val[I];
						lrs.Sbx += val[G]*val[P];
						lrs.Sby += val[R]*val[P];
						lrs.Sbz += val[I]*val[P];
						lrs.Sx += val[G];
						lrs.Sy += val[R];
						lrs.Sz += val[I];
						lrs.Sb += val[P];
					}
				break;

				case AC_STATS_1:
					if (val[0] || val[1] || val[2])
					{
						calculate_pixel(p, val, stats, 0);
						// line_S += 0.299*p[R] + 0.587*p[G] + 0.114*p[B];
						line_S += p[R] + p[G] + p[B];
						if (val[R] && val[G] && val[B]) pixel_count += 3;
						for (int i = 0; i < 3; ++i)
						{
							if (val[i] == 0) continue;

							ls[i].Sx += p[i];
							++ls[i].pixel_count;
						}
					}
				break;

				case AC_STATS_2:
					if (val[0] || val[1] || val[2])
					{
						calculate_pixel(p, val, stats, 0);
						// line_S += sqr(0.299*p[R] + 0.587*p[G] + 0.114*p[B] - avg_Y);
						line_S += sqr(p[R] - avg_Y);
						line_S += sqr(p[G] - avg_Y);
						line_S += sqr(p[B] - avg_Y);
						for (int i = 0; i < 3; ++i)
						{
							if (val[i] == 0) continue;

							ls[i].Sdd += sqr(p[i] - stats[i].avg);
						}
					}
				break;

				default:
					; // Do nothing
				break;
			}

			// Move on...
			pixel += 3;
		}

		switch (action)
		{
			case AC_MERGE:
				// Write the cropped portion of the scanline
				fwrite(outline + crop_x*3, crop_w*3, 1, out_rgb);
			break;

			case AC_STATS_1:
				Sy += line_S;
				for (int i = 0; i < 3; ++i)
				{
					stats[i].Sx += ls[i].Sx;
					stats[i].pixel_count += ls[i].pixel_count;
				}
			break;

			case AC_STATS_2:
				Sdd += line_S;
				for (int i = 0; i < 3; ++i)
					stats[i].Sdd += ls[i].Sdd;
			break;

			case AC_REGRESSION:
				add_regr_stats(&regr_stats, &lrs);
			break;

			default:
				;
			break;
		}

		// Print some progress
		if (y % 256 == 0) progress("Line %d processed.\r", y);
	}
	printf("calculation done.                    \n");

	// Free memory
	free(outline);
	for (int col = 0; col < 5; ++col)
		free(current[col]);
}

// Calculates alpha, beta, gamma
void calc_abg()
{
	double sxx = regr_stats.Sxx, syy = regr_stats.Syy, szz = regr_stats.Szz,
	       sxy = regr_stats.Sxy, syz = regr_stats.Syz, sxz = regr_stats.Sxz,
	       sbx = regr_stats.Sbx, sby = regr_stats.Sby, sbz = regr_stats.Sbz,
	       sx = regr_stats.Sx, sy = regr_stats.Sy, sz = regr_stats.Sz, sb = regr_stats.Sb;

	double alpha1 = sb*sx*syz*syz - sb*sx*syy*szz + sb*sxy*sy*szz
		- sb*sxy*syz*sz - sb*sxz*sy*syz + sb*sxz*syy*sz - sbx*sy*sy*szz
		+ 2*sbx*sy*syz*sz - sbx*syy*sz*sz + sbx*syy*szz - sbx*syz*syz
		+ sby*sx*sy*szz - sby*sx*syz*sz + sby*sxy*sz*sz - sby*sxy*szz
		- sby*sxz*sy*sz + sby*sxz*syz - sbz*sx*sy*syz + sbz*sx*syy*sz
		- sbz*sxy*sy*sz + sbz*sxy*syz + sbz*sxz*sy*sy - sbz*sxz*syy;
	double beta1 = sb*sx*sxy*szz - sb*sx*sxz*syz - sb*sxx*sy*szz
		+ sb*sxx*syz*sz - sb*sxy*sxz*sz + sb*sxz*sxz*sy + sbx*sx*sy*szz
		- sbx*sx*syz*sz + sbx*sxy*sz*sz - sbx*sxy*szz - sbx*sxz*sy*sz
		+ sbx*sxz*syz - sby*sx*sx*szz + 2*sby*sx*sxz*sz - sby*sxx*sz*sz
		+ sby*sxx*szz - sby*sxz*sxz + sbz*sx*sx*syz - sbz*sx*sxy*sz
		- sbz*sx*sxz*sy + sbz*sxx*sy*sz - sbz*sxx*syz + sbz*sxy*sxz;
	double gamma1 = sb*sx*sxz*syy - sb*sx*sxy*syz + sb*sxx*sy*syz
		- sb*sxx*syy*sz + sb*sxy*sxy*sz - sb*sxy*sxz*sy - sbx*sx*sy*syz
		+ sbx*sx*syy*sz - sbx*sxy*sy*sz + sbx*sxy*syz + sbx*sxz*sy*sy
		- sbx*sxz*syy + sby*sx*sx*syz - sby*sx*sxy*sz - sby*sx*sxz*sy
		+ sby*sxx*sy*sz - sby*sxx*syz + sby*sxy*sxz - sbz*sx*sx*syy
		+ 2*sbz*sx*sxy*sy - sbz*sxx*sy*sy + sbz*sxx*syy - sbz*sxy*sxy;
	double phi1 = sb*sxx*syy*szz - sb*sxx*syz*syz - sb*sxy*sxy*szz + 2*sb*sxy*sxz*syz
		- sb*sxz*sxz*syy - sbx*sx*syy*szz + sbx*sx*syz*syz + sbx*sxy*sy*szz
		- sbx*sxy*syz*sz - sbx*sxz*sy*syz + sbx*sxz*syy*sz + sby*sx*sxy*szz
		- sby*sx*sxz*syz - sby*sxx*sy*szz + sby*sxx*syz*sz - sby*sxy*sxz*sz
		+ sby*sxz*sxz*sy - sbz*sx*sxy*syz + sbz*sx*sxz*syy + sbz*sxx*sy*syz
		- sbz*sxx*syy*sz + sbz*sxy*sxy*sz - sbz*sxy*sxz*sy;
	double d = sx*sx*syz*syz - sx*sx*syy*szz + 2*sx*sxy*sy*szz - 2*sx*sxy*syz*sz
		- 2*sx*sxz*sy*syz + 2*sx*sxz*syy*sz - sxx*sy*sy*szz
		+ 2*sxx*sy*syz*sz - sxx*syy*sz*sz + sxx*syy*szz - sxx*syz*syz
		+ sxy*sxy*sz*sz - sxy*sxy*szz - 2*sxy*sxz*sy*sz + 2*sxy*sxz*syz
		+ sxz*sxz*sy*sy - sxz*sxz*syy;

	alpha = alpha1 / d;
	beta = beta1 / d;
	gamma = gamma1 / d;
	phi = phi1 / d;

	printf("using panchro composition weight vector: %g, %g, %g, %g\n", alpha, beta, gamma, phi);
}

/// Prints usage info & quits
void usage()
{
	printf("usage:\n");
	printf("\t./merge red.tif green.tif blue.tif near_ir.tif panchrom.tif <crop_definition> output_file.rgb eta <scene_rect>\n");
	printf("crop definition:\n\tWIDTHxHEIGHT@X,Y  e.g. 800x600@1017,534\n");
	exit(1);
}

/// Parses the crop string (WIDTHxHEIGHT@X,Y)
void parse_crop(char *str)
{
	char *x = strchr(str, 'x');
	char *at = strchr(str, '@');
	char *comma = strchr(str, ',');

	if (!x || !at || !comma)
	{
		printf("Invalid crop format.\n");
		usage();
	}

	*x = *at = *comma = '\0';
	crop_w = atoi(str);
	crop_h = atoi(x+1);
	crop_x = atoi(at+1);
	crop_y = atoi(comma+1);
}

int main(int argc, char *argv[])
{

	if (argc != 18) usage();

	// Parse options
	parse_crop(argv[6]);
	const char *outfn = argv[7];
	eta = atof(argv[8]);
	psi = atof(argv[9]);
	tlx = atof(argv[10]);
	tly = atof(argv[11]);
	trx = atof(argv[12]);
	try = atof(argv[13]);
	blx = atof(argv[14]);
	bly = atof(argv[15]);
	brx = atof(argv[16]);
	bry = atof(argv[17]);

	// Open the files
	tiff_in_t tiffs;
	for (int col = 0; col < 5; ++col)
	{
		tiffs.chan[col] = tiff_open_r(argv[col+1]);
		if (!tiffs.chan[col])
		{
			die("Could not open channel: %s!\n", argv[col+1]);
		}
	}
	FILE *out_rgb = fopen(outfn, "w");
	if (!out_rgb) die("%s: Could not open for writing.\n", outfn);

	// Calculate regression coefficients
	regr_stats = zero_regr_stats;
	process_images(&tiffs, AC_REGRESSION, stats, out_rgb);
	calc_abg();

	// Calculate RGB stats for postprocess normalization
	for (int i = 0; i < 3; ++i) stats[i] = zero_stats;

	process_images(&tiffs, AC_STATS_1, stats, out_rgb);
	printf("average: ");
	for (int i = 0; i < 3; ++i) 
	{
		stats[i].avg = (stats[i].Sx / stats[i].pixel_count);
		printf("%g ", stats[i].avg);
	}
	printf("\n");
	avg_Y = Sy / pixel_count;
	printf("average Y: %g\n", avg_Y);

	process_images(&tiffs, AC_STATS_2, stats, out_rgb);
	printf("standard deviation: ");
	for (int i = 0; i < 3; ++i)
	{
		stats[i].sd = sqrt(stats[i].Sdd / stats[i].pixel_count);
		printf("%g ", stats[i].sd);
	}
	printf("\n");
	sd_Y = sqrt(Sdd / pixel_count);
	printf("Y standard deviation: %g\n", sd_Y);

	// Merge the images
	process_images(&tiffs, AC_MERGE, stats, out_rgb);

	// Close the files
	for (int col = 0; col < 5; ++col)
		tiff_close(tiffs.chan[col]);
	fclose(out_rgb);

	return 0;
}
